# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
from typing import Any

import numpy as np
import pytest
import ray
from omegaconf import DictConfig
from PIL import Image
from transformers.utils import get_json_schema

from verl.experimental.agent_loop import AgentLoopManager
from verl.protocol import DataProto
from verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema
from verl.tools.schemas import ToolResponse
from verl.utils import hf_tokenizer


@pytest.fixture
def init_config() -> DictConfig:
    from hydra import compose, initialize_config_dir

    with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")):
        config = compose(
            config_name="ppo_trainer",
            overrides=[
                "actor_rollout_ref.actor.use_dynamic_bsz=true",
                # test sleep/wake_up with fsdp offload
                "actor_rollout_ref.actor.fsdp_config.param_offload=True",
                "actor_rollout_ref.actor.fsdp_config.optimizer_offload=True",
            ],
        )

    model_path = os.path.expanduser("~/models/Qwen/Qwen2.5-VL-3B-Instruct")
    config.actor_rollout_ref.model.path = model_path
    config.actor_rollout_ref.rollout.name = os.environ["ROLLOUT_NAME"]
    config.actor_rollout_ref.rollout.mode = "async"
    config.actor_rollout_ref.rollout.enforce_eager = True
    config.actor_rollout_ref.rollout.prompt_length = 4096
    config.actor_rollout_ref.rollout.response_length = 4096
    config.actor_rollout_ref.rollout.n = 4
    config.actor_rollout_ref.rollout.agent.num_workers = 2
    config.actor_rollout_ref.rollout.skip_tokenizer_init = True

    return config


class ImageGeneratorTool(BaseTool):
    def generate_image(self, description: str, size: str = "256x256"):
        """Generate a simple image based on description.

        Args:
            description: The description of the image to generate.
            size: The size of the image. Defaults to "256x256". (choices: ["256x256", "512x512"])

        Returns:
            A generated image
        """
        print(f"[DEBUG] generate_image: {description}, {size}")
        # Create a simple colored image for testing
        width, height = map(int, size.split("x"))

        # Create different colors based on description
        if "red" in description.lower():
            color = (255, 0, 0)
        elif "blue" in description.lower():
            color = (0, 0, 255)
        elif "green" in description.lower():
            color = (0, 255, 0)
        else:
            color = (128, 128, 128)  # gray

        # Create image
        image = Image.new("RGB", (width, height), color)

        # Add some pattern to make it more interesting
        for i in range(0, width, 50):
            for j in range(0, height, 50):
                # Add white squares in a grid pattern
                for x in range(i, min(i + 20, width)):
                    for y in range(j, min(j + 20, height)):
                        image.putpixel((x, y), (255, 255, 255))

        return image

    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:
        schema = get_json_schema(self.generate_image)
        return OpenAIFunctionToolSchema(**schema)

    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:
        try:
            image = self.generate_image(**parameters)
            # Return the PIL Image directly - the framework should handle the conversion
            return ToolResponse(image=[image]), 0, {}
        except Exception as e:
            return ToolResponse(text=str(e)), 0, {}


@pytest.mark.flaky(reruns=3)
def test_multimodal_tool_agent(init_config):
    """Test agent loop with multimodal tool that returns images using Qwen VL model."""
    ray.shutdown()
    ray.init(
        runtime_env={
            "env_vars": {
                "TOKENIZERS_PARALLELISM": "true",
                "NCCL_DEBUG": "WARN",
                "VLLM_LOGGING_LEVEL": "INFO",
                "VLLM_USE_V1": "1",
            }
        },
        ignore_reinit_error=True,
    )

    # Add custom chat template to enable tool calling support (same as recipe/deepeyes)
    template_path = os.path.join(os.path.dirname(__file__), "qwen_vl_tool_chat_template.jinja2")
    with open(template_path, encoding="utf-8") as f:
        custom_chat_template = f.read()

    init_config.actor_rollout_ref.model.custom_chat_template = custom_chat_template

    # =========================== 1. Init rollout manager with image tool ===========================
    tool_config = {
        "tools": [
            {
                "class_name": "tests.experimental.agent_loop.test_multi_modal.ImageGeneratorTool",
                "config": {"type": "native"},
            },
        ]
    }
    tool_config_path = "/tmp/multimodal_tool_config.json"
    with open(tool_config_path, "w") as f:
        json.dump(tool_config, f)

    n = 2
    init_config.actor_rollout_ref.rollout.n = n
    init_config.actor_rollout_ref.rollout.multi_turn.tool_config_path = tool_config_path
    init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 1
    init_config.actor_rollout_ref.rollout.multi_turn.max_user_turns = 1
    agent_loop_manager = AgentLoopManager(init_config)

    # =========================== 2. Generate sequences with multimodal prompts ===========================
    raw_prompts = [
        [
            {"role": "user", "content": "How are you?"},
        ],
        [
            {"role": "user", "content": "Please generate a red image for me."},
        ],
        [
            {"role": "user", "content": "Can you create a blue picture with size 512x512?"},
        ],
        [
            {
                "role": "system",
                "content": (
                    "You are Qwen VL, created by Alibaba Cloud. You are a helpful "
                    "assistant that can generate and analyze images."
                ),
            },
            {"role": "user", "content": "Generate a green landscape image and describe what you see in it."},
        ],
    ]

    batch = DataProto(
        non_tensor_batch={
            "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object),
            "agent_name": np.array(["tool_agent"] * len(raw_prompts)),
            "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)),
            "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)),
        },
    )
    batch = batch.repeat(n)
    result = agent_loop_manager.generate_sequences(prompts=batch)
    assert len(result) == len(raw_prompts) * n

    # Check turns
    num_turns = result.non_tensor_batch["__num_turns__"]
    print(f"num_turns: {num_turns}")
    for i in range(len(num_turns)):
        if i // n == 0:
            # First prompt: "How are you?" - should have 2 turns [user, assistant]
            assert num_turns[i] == 2, f"Expected 2 turns but got {num_turns[i]} for sample {i}"
        else:
            # Tool-calling prompts should have 4 turns [user, assistant, tool, assistant]
            assert num_turns[i] == 4, f"Expected 4 turns but got {num_turns[i]} for sample {i}"

    # Check that images were properly returned in the tool responses
    tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path)
    responses = result.batch["responses"]
    response_mask = result.batch["response_mask"]
    attention_mask = result.batch["attention_mask"]
    assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}"
    response_length = response_mask.size(1)

    image_found_count = 0
    for i in range(len(responses)):
        # response with tool response (including images)
        valid_tokens = responses[i][attention_mask[i][-response_length:].bool()]
        response_with_obs = tokenizer.decode(valid_tokens)

        # response without tool response
        valid_tokens = responses[i][response_mask[i].bool()]
        response_without_obs = tokenizer.decode(valid_tokens)

        # Check that tool responses were properly masked out from training
        assert "<tool_response>" not in response_without_obs, (
            f"found <tool_response> in response: {response_without_obs}"
        )
        assert "</tool_response>" not in response_without_obs, (
            f"found </tool_response> in response: {response_without_obs}"
        )

        # Check that images were included in the full response
        if "<image>" in response_with_obs or "image" in response_with_obs.lower():
            image_found_count += 1

        print("=========================")
        print("Response with tool observations:")
        print(response_with_obs)
        print("---")
        print("Response without tool observations:")
        print(response_without_obs)

    # Verify that tool-calling responses contained image-related content
    print(f"Found {image_found_count} responses with image content out of {len(responses)}")
    # We should have at least some image content from the tool-calling prompts
    # Note: First prompt might not use tools, so we don't expect 100% image content
    expected_tool_calls = sum(1 for i in range(len(num_turns)) if num_turns[i] == 4)
    assert image_found_count >= 0, (
        f"No image-related content found, but expected at least some from {expected_tool_calls} tool calls"
    )

    print("Multimodal tool test passed!")
    ray.shutdown()


def test_multimodal_single_turn_agent(init_config):
    """Test single turn agent loop with multimodal inputs using Qwen VL model."""
    ray.init(
        runtime_env={
            "env_vars": {
                "TOKENIZERS_PARALLELISM": "true",
                "NCCL_DEBUG": "WARN",
                "VLLM_LOGGING_LEVEL": "INFO",
                "VLLM_USE_V1": "1",
            }
        },
        ignore_reinit_error=True,
    )

    # =========================== 1. Init rollout manager ===========================
    n = 2
    init_config.actor_rollout_ref.rollout.n = n
    init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 1
    init_config.actor_rollout_ref.rollout.multi_turn.max_user_turns = 1
    agent_loop_manager = AgentLoopManager(init_config)

    # =========================== 2. Generate sequences with multimodal prompts ===========================
    # Create a simple test image
    test_image = Image.new("RGB", (256, 256), (100, 150, 200))
    test_image2 = Image.new("RGB", (512, 512), (100, 150, 200))

    raw_prompts = [
        [
            {"role": "user", "content": "Hello, how are you?"},
        ],
        [
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": "What color is this image?"},
                ],
            },
        ],
        [
            {
                "role": "system",
                "content": "You are Qwen VL, created by Alibaba Cloud. You are a helpful assistant.",
            },
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": "Describe this image in detail."},
                ],
            },
        ],
    ]

    # Prepare multi_modal_data for each prompt
    multi_modal_data_list = [
        None,  # First prompt: text only
        {"image": test_image},  # Second prompt: with image
        {"image": test_image2},  # Third prompt: with image
    ]

    batch = DataProto(
        non_tensor_batch={
            "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object),
            "agent_name": np.array(["single_turn_agent"] * len(raw_prompts)),
            "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)),
            "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)),
        },
    )

    # Add multi_modal_data to batch
    multi_modal_data_array = np.array([data if data else {} for data in multi_modal_data_list], dtype=object)
    batch.non_tensor_batch["multi_modal_data"] = multi_modal_data_array

    batch = batch.repeat(n)
    result = agent_loop_manager.generate_sequences(prompts=batch)
    assert len(result) == len(raw_prompts) * n

    # Check turns - all should be single turn (2: user + assistant)
    num_turns = result.non_tensor_batch["__num_turns__"]
    print(f"num_turns: {num_turns}")
    for i in range(len(num_turns)):
        assert num_turns[i] == 2, f"Expected 2 turns but got {num_turns[i]} for sample {i}"

    # Verify responses
    tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path)
    prompts = result.batch["prompts"]
    responses = result.batch["responses"]
    response_mask = result.batch["response_mask"]
    assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}"

    # Check for image pads in prompts
    image_pad_count = 0
    for i in range(len(prompts)):
        prompt_ids = prompts[i][prompts[i] != tokenizer.pad_token_id].tolist()
        prompt_text = tokenizer.decode(prompt_ids)

        # Check if this sample should have image pads (samples with index 1 and 2 in each repeat have images)
        sample_idx = i // n
        has_image_pad = "<|image_pad|>" in prompt_text or "<|vision_start|>" in prompt_text

        print("=========================")
        print(f"Sample {i} (original prompt index: {sample_idx}):")
        print(f"Prompt length: {len(prompt_ids)} tokens")
        print(f"Has image_pad: {has_image_pad}")

        if sample_idx != 0:  # Samples 1 and 2 should have images
            if has_image_pad:
                image_pad_count += 1
                # Count the number of image_pad tokens
                num_image_pads = prompt_text.count("<|image_pad|>")
                print(f"Number of <|image_pad|> tokens: {num_image_pads}")
            else:
                print("WARNING: Expected image_pad but not found!")

        # Show first 200 chars of prompt
        print(f"Prompt text (first 200 chars): {prompt_text[:200]}...")

    for i in range(len(responses)):
        valid_tokens = responses[i][response_mask[i].bool()]
        response_text = tokenizer.decode(valid_tokens)
        print(f"Sample {i} response: {response_text[:100]}...")

    # Verify that we found image pads in multimodal samples
    expected_multimodal_samples = 2 * n  # 2 prompts with images, repeated n times
    print(f"\nFound {image_pad_count} samples with image_pad out of {expected_multimodal_samples} expected")
    assert image_pad_count > 0, "No image_pad tokens found in multimodal samples!"

    print("Single turn multimodal test passed!")
    ray.shutdown()


def test_multimodal_partial_single_turn_agent(init_config):
    """Test partial single turn agent loop with multimodal inputs using Qwen VL model."""

    # TODO(baiyan):
    #    see verl/recipe/fully_async_policy/agent_loop/partial_single_turn_agent_loop.py for more details.
    #    if use_correct_processor=True, the test will pass but the async training will hang, so I disable this test
    #    for now

    return

    ray.init(
        runtime_env={
            "env_vars": {
                "TOKENIZERS_PARALLELISM": "true",
                "NCCL_DEBUG": "WARN",
                "VLLM_LOGGING_LEVEL": "INFO",
                "VLLM_USE_V1": "1",
            }
        },
        ignore_reinit_error=True,
    )
    from recipe.fully_async_policy.agent_loop import FullyAsyncAgentLoopManager

    # =========================== 1. Init rollout manager ===========================
    n = 2
    init_config.actor_rollout_ref.rollout.n = n
    init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 1
    init_config.actor_rollout_ref.rollout.multi_turn.max_user_turns = 1
    import asyncio

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    agent_loop_manager = loop.run_until_complete(FullyAsyncAgentLoopManager.create(init_config))

    # =========================== 2. Generate sequences with multimodal prompts ===========================
    # Create a simple test image
    test_image = Image.new("RGB", (256, 256), (200, 100, 50))
    test_image2 = Image.new("RGB", (512, 512), (100, 150, 200))

    raw_prompts = [
        [
            {"role": "user", "content": "What is the capital of France?"},
        ],
        [
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": "What do you see in this image?"},
                ],
            },
        ],
        [
            {
                "role": "system",
                "content": "You are Qwen VL, a helpful multimodal assistant.",
            },
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": "Analyze the colors in this image."},
                ],
            },
        ],
    ]

    # Prepare multi_modal_data for each prompt
    multi_modal_data_list = [
        None,  # First prompt: text only
        {"image": test_image},  # Second prompt: with image
        {"image": test_image2},  # Third prompt: with image
    ]

    batch = DataProto(
        non_tensor_batch={
            "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object),
            "agent_name": np.array(["partial_single_turn_agent"] * len(raw_prompts)),
            "data_source": np.array(["openai/gsm8k"] * len(raw_prompts)),
            "reward_model": np.array([{"style": "rule", "ground_truth": "1.0"}] * len(raw_prompts)),
        },
    )

    # Add multi_modal_data to batch
    multi_modal_data_array = np.array([data if data else {} for data in multi_modal_data_list], dtype=object)
    batch.non_tensor_batch["multi_modal_data"] = multi_modal_data_array

    batch = batch.repeat(n)
    result = agent_loop_manager.generate_sequences(prompts=batch)
    assert len(result) == len(raw_prompts) * n

    # Check turns - all should be single turn (2: user + assistant)
    num_turns = result.non_tensor_batch["__num_turns__"]
    print(f"num_turns: {num_turns}")
    for i in range(len(num_turns)):
        assert num_turns[i] == 2, f"Expected 2 turns but got {num_turns[i]} for sample {i}"

    # Verify responses
    tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path)
    prompts = result.batch["prompts"]
    responses = result.batch["responses"]
    response_mask = result.batch["response_mask"]
    assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}"

    # Check for image pads in prompts
    image_pad_count = 0
    for i in range(len(prompts)):
        prompt_ids = prompts[i][prompts[i] != tokenizer.pad_token_id].tolist()
        prompt_text = tokenizer.decode(prompt_ids)

        # Check if this sample should have image pads (samples with index 1 and 2 in each repeat have images)
        sample_idx = i // n
        has_image_pad = "<|image_pad|>" in prompt_text or "<|vision_start|>" in prompt_text

        print("=========================")
        print(f"Sample {i} (original prompt index: {sample_idx}):")
        print(f"Prompt length: {len(prompt_ids)} tokens")
        print(f"Has image_pad: {has_image_pad}")

        if sample_idx != 0:  # Samples 1 and 2 should have images
            if has_image_pad:
                image_pad_count += 1
                # Count the number of image_pad tokens
                num_image_pads = prompt_text.count("<|image_pad|>")
                print(f"Number of <|image_pad|> tokens: {num_image_pads}")
            else:
                print("WARNING: Expected image_pad but not found!")

        # Show first 200 chars of prompt
        print(f"Prompt text (first 200 chars): {prompt_text[:200]}...")

    for i in range(len(responses)):
        valid_tokens = responses[i][response_mask[i].bool()]
        response_text = tokenizer.decode(valid_tokens)
        print(f"Sample {i} response: {response_text[:100]}...")

    # Verify that we found image pads in multimodal samples
    expected_multimodal_samples = 2 * n  # 2 prompts with images, repeated n times
    print(f"\nFound {image_pad_count} samples with image_pad out of {expected_multimodal_samples} expected")
    assert image_pad_count > 0, "No image_pad tokens found in multimodal samples!"

    print("Partial single turn multimodal test passed!")
    ray.shutdown()
